import numpy as np
import pyopencl as cl
from pyopencl.array import Array
from verbose import Verbose

W_P_PAR = 4 #width of pair parameters, float-4,8,16
W_B_PAR = 4 #width of bond parameters, float-4,8,16
W_A_PAR = 4 #width of angle parameters, float-4,8,16

N_P_NEIGHBOR = 128  #number of pair neighbors, stored in int16*, thus must be 16*N 

echo = Verbose().echo
class Configuration:
    """
    A configuration of POCLMD includes normal configuration and static topology.
    The information is usually needed by all other parts of POCLMD.
    ---host---
    float4: box_lo, box_hi, box_l, box_p,
    int: n_atoms, n_bonds, n_angles,
    int: n_atom_type, n_bond_type, n_angle_type,
    float4*: r_t, v_w, f_e, [xyz_type,velocity_1/m,force_en]
    float4*: rr_c [reduced-xyz_cell-index]
    float8*: vir, [virial]
    int2*: bond
    int3*: angle
    ---device/gpu---
    int: _n_atom_type, _n_bond_type, _n_angle_type,
    float4*: _r_t, _v_w, _f_e, [xyz_type,velocity_1/m,force_en]
    float8*: _vir, [virial]
    int*: _excluded_neighbor
    int*: _bond_neighbor,
    int*: _anlge0_neighbor,_anlge1_neighbor,_anlge2_neighbor
    ---reduced properties---
    ?float*: _tot_e(E), _tot_vir(P), _tot_v2(T)
    =====
    on opencl-device: [name] --> _[name]
    eg. r_t (on host) --> _r_t (on device/gpu)
    =====
    """
    def __init__(self):
        pass

    def init_device(self,context):
        mf = cl.mem_flags
        nbytes = self.n_atoms * 16 #float4 = 16 bytes
        print 'float:',nbytes, self.r_t.nbytes
        self._r_t = cl.Buffer(context,mf.READ_WRITE, size=nbytes)
        self._v_w = cl.Buffer(context,mf.READ_WRITE, size=nbytes)
        self._f_e = cl.Buffer(context,mf.READ_WRITE, size=nbytes)
        self._rr_c = cl.Buffer(context,mf.READ_WRITE, size=nbytes)
        nbytes = self.n_atoms * 32 #float8 = 32 bytes
        self._vir = cl.Buffer(context,mf.READ_WRITE, size=nbytes)
        #### self._bond_neighbor ... = 
        echo('Configuration: init_device done',10) #debug
        #
        
    def copy_to_device(self,queue):
        cl.enqueue_copy(queue, self._r_t, self.r_t) #dest <-- src
        cl.enqueue_copy(queue, self._v_w, self.v_w) #dest <-- src
        #### self._bond_neighbor ... 
        echo('Configuration: copy_to_device done',10) #debug
        #for test:
        cl.enqueue_copy(queue, self._f_e, self.f_e) #dest <-- src

    def get_from_device(self,queue):
        cl.enqueue_copy(queue, self.r_t, self._r_t) #dest <-- src
        cl.enqueue_copy(queue, self.v_w, self._v_w) #dest <-- src
        cl.enqueue_copy(queue, self.f_e, self._f_e) #dest <-- src
        cl.enqueue_copy(queue, self.vir, self._vir) #dest <-- src
        echo('Configuration: get_from_device done',10) #debug

    def get_bond_neighbor_list(self,bonds):
        echo('Configuration: get_bond_neighbor_list done',10) #debug

    def dget_angle_neighbor_list(self,angles):
        echo('Configuration: get_angle_neighbor_list done',10) #debug

    def get_excluded_neighbor_list(self,bonds,level=3):
        """
        level- 
        1:upto bonded; o-o
        2:upto angled: o-o-o
        3:upto dihedrial: o-o-o-o
        """
        echo('Configuration: get_excluded_neighbor_list done',10) #debug

    #-----
    def read_lammps_data(self,datafile,datatype='pair'):
        """
        read lammps data file
        """
        header_keywords = [ "atoms","bonds","angles","dihedrals","impropers", 
                            "atom types","bond types","angle types",
                            "dihedral types","improper types",
                            "xlo xhi","ylo yhi","zlo zhi","xy xz yz" ] 
        section_keywords = [ "Masses","Atoms","Velocities",
                             "Bonds","Angles", 
                             "Dihedrals","Impropers"]

        if(datafile[-3:]=='.gz'):
            import gzip
            f = gzip.open(datafile,'r')
        else:
            f = open(datafile,'r')
        self.data_title = f.readline()
        self.box_lo = np.zeros((4,),'float32')
        self.box_hi = np.zeros((4,),'float32')
        #read header
        while True:
            line = f.readline().strip()
            if len(line) == 0:
                continue #skip blank line
            found = False
            for keyword in header_keywords:
                if line.find(keyword) >= 0:
                    found = True
                    words = line.split()
                    # --
                    if   keyword == "atoms":
                        self.n_atoms = int(words[0])
                    elif keyword == "bonds":
                        self.n_bonds = int(words[0])
                    elif keyword == "angles":
                        self.n_angles = int(words[0])
                    elif keyword == "dihedrals":
                        self.n_dihedrals = int(words[0])
                    elif keyword == "impropers":
                        self.n_impropers = int(words[0])
                    # --
                    elif keyword == "atom types":
                        self.n_atom_type = int(words[0])
                    elif keyword == "bond types":
                        self.n_bond_type = int(words[0])
                    elif keyword == "angle types":
                        self.n_angle_type = int(words[0])
                    elif keyword == "dihedral types":
                        self.n_dihedral_type = int(words[0])
                    elif keyword == "improper types":
                        self.n_improper_type = int(words[0])
                    # --
                    elif keyword == "xlo xhi":
                        self.box_lo[0] = float(words[0])
                        self.box_hi[0] = float(words[1])
                    elif keyword == "ylo yhi":
                        self.box_lo[1] = float(words[0])
                        self.box_hi[1] = float(words[1])
                    elif keyword == "zlo zhi":
                        self.box_lo[2] = float(words[0])
                        self.box_hi[2] = float(words[1])
            if not found: #begining of data-section
                    break
        echo('Configuration: read_lammps_data read header of %s done.'%(datafile),10) #debug
        #pre-allocation and additional misc:
        self.box_l = self.box_hi - self.box_lo #box-size
        self.box_l[3] = 1. #avoid divided by 0
        self.box_p = np.array([1,1,1,0],'float32') #periodic
        self.mass = np.ones((self.n_atom_type,),'float32')
        self.r_t = np.zeros((self.n_atoms,4),'float32') #xyz_type
        self.img = np.zeros((self.n_atoms,4),'float32') #xyz-image
        self.v_w = np.zeros((self.n_atoms,4),'float32') #velocity_1/mass
        self.f_e = np.zeros((self.n_atoms,4),'float32') #force_en
        #--
        self.vir = np.zeros((self.n_atoms,8),'float32') #viral
        #--
        #read sections
        while True:
            found = False
            for keyword in section_keywords:
                if keyword == line:
                    found = True
                    f.readline() #there is a blank line between keyword and data
                    if   keyword == "Masses": #type,mass
                        for i in range(self.n_atom_type):
                            s = f.readline().strip().split()
                            it = int(s[0]) - 1 #lammps index from 1 not 0
                            self.mass[it] = float(s[1])
                    elif keyword == "Atoms": #id,type,x,y,z,ix,iy,iz
                        if   datatype == 'pair':
                            for i in range(self.n_atoms):
                                s = f.readline().strip().split()
                                ida = int(s[0]) - 1 #id, lammps index from 1 not 0
                                ita = int(s[1]) - 1 #tpye
                                self.r_t[i,0] = float(s[2])
                                self.r_t[i,1] = float(s[3])
                                self.r_t[i,2] = float(s[4])
                                self.r_t[i,3] = float(ita) #type
                                self.img[i,0] = float(s[5])
                                self.img[i,1] = float(s[6])
                                self.img[i,2] = float(s[7])
                        elif datatype == 'molecular': #id,id-mol,type,x,y,z,ix,iy,iz
                            for i in range(self.n_atoms):
                                s = f.readline().strip().split()
                                ida = int(s[0]) - 1 #id, lammps index from 1 not 0
                                idm = int(s[1]) - 1 #id-molecular
                                ita = int(s[2]) - 1 #type
                                self.r_t[i,0] = float(s[3])
                                self.r_t[i,1] = float(s[4])
                                self.r_t[i,2] = float(s[5])
                                self.r_t[i,3] = float(ita) #type
                                self.img[i,0] = float(s[6])
                                self.img[i,1] = float(s[7])
                                self.img[i,2] = float(s[8])
                                self.img[i,3] = 0.
                    elif keyword == "Velocities": #id,vx-vy-vz
                        for i in range(self.n_atoms):
                            s = f.readline().strip().split()
                            ida = int(s[0]) - 1 #id, lammps index from 1 not 0
                            self.v_w[i,0] = float(s[1])
                            self.v_w[i,1] = float(s[2])
                            self.v_w[i,2] = float(s[3])
                    ## -elif keyword == "Bonds": 
                    #--
            if not found:
                raise StandardError, "invalid section %s in data file" % line
            f.readline() #there is a blank line between data and next keyword
            line = f.readline()
            if not line: #readline() return empty at EOF
                break
            line = line.strip()
        f.close()
        echo('Configuration: read_lammps_data read sections of %s done.'%(datafile),10) #debug
        #post-procedure
        self.r_t += self.box_lo + self.img*self.box_l #restore real value
        for i in range(self.n_atoms):
            it = self.r_t[i,3]
            self.v_w[i,3] = 1./self.mass[it] #inverse of mass
        #--bond,angle,exclude -lists are static, done by CPU once
        #self.bond = None
        #self.angle = None
        echo('Configuration: read_lammps_data post-procedure done.',10) #debug

